#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
扩展TextArea控件，支持ANSI颜色和原生滚动功能
"""
import logging
from prompt_toolkit.document import Document
from prompt_toolkit.layout.containers import Window
from prompt_toolkit.layout.controls import BufferControl
from prompt_toolkit.widgets import TextArea
from prompt_toolkit.formatted_text import ANSI, to_formatted_text
from prompt_toolkit.data_structures import Point
from prompt_toolkit.layout.controls import UIContent
from prompt_toolkit.filters import Condition, to_filter
from prompt_toolkit.layout.processors import (
    AppendAutoSuggestion,
    BeforeInput,
    ConditionalProcessor,
    PasswordProcessor,
)
from prompt_toolkit.filters import has_focus, is_done
from prompt_toolkit.lexers import DynamicLexer
from prompt_toolkit.mouse_events import MouseEvent


class AnsiBufferControl(BufferControl):
    """
    支持ANSI颜色代码的BufferControl
    """
    def __init__(self, buffer, **kwargs):
        super().__init__(buffer=buffer, **kwargs)
    
    def create_content(self, width: int, height: int) -> UIContent:
        """
        生成此用户控件的内容，支持ANSI颜色代码
        """
        buffer = self.buffer
        document = buffer.document
        
        if not buffer or not buffer.text:
            return UIContent(
                get_line=lambda i: [],
                line_count=0,
                cursor_position=Point(0, 0)
            )

        def get_line(line_number: int):
            # 获取指定行的文本
            lines = buffer.text.split('\n')
            if line_number >= len(lines):
                return []
            
            line = lines[line_number]
            # 使用ANSI处理颜色代码
            ansi_text = ANSI(line)
            # 转换为格式化文本
            fragments = to_formatted_text(ansi_text)
            # 添加一个空格，这是光标可能的位置（原生BufferControl也会这样做）
            fragments = fragments + [("", " ")]
            return fragments

        # 使用与原生BufferControl相同的光标位置计算方法
        cursor_position = Point(x=document.cursor_position_col, y=document.cursor_position_row)
        
        return UIContent(
            get_line=get_line,
            line_count=document.line_count,
            cursor_position=cursor_position,
        )
        
    def mouse_handler(self, mouse_event: MouseEvent):
        """
        处理鼠标事件，对于不处理的事件返回NotImplemented
        """
        # 对于我们不处理的鼠标事件，返回NotImplemented，让Window处理
        return NotImplemented


class AnsiTextArea(TextArea):
    """
    支持ANSI颜色和原生滚动的TextArea
    """
    
    def __init__(self, auto_scroll_bottom=True, **kwargs):
        # 设置为只读模式
        kwargs['read_only'] = True
        
        # 先调用父类构造函数，创建所有组件
        super().__init__(**kwargs)

        # 保存参数
        preview_search = kwargs.get('preview_search', True)
        focusable = kwargs.get('focusable', True)
        focus_on_click = kwargs.get('focus_on_click', True)
        password = kwargs.get('password', False)
        prompt = kwargs.get('prompt', '')
        input_processors = kwargs.get('input_processors', [])

        self.auto_scroll_bottom = auto_scroll_bottom
        
        # 创建支持ANSI的BufferControl，替换原有的control
        self.control = AnsiBufferControl(
            buffer=self.buffer,
            lexer=DynamicLexer(lambda: self.lexer),
            input_processors=[
                ConditionalProcessor(
                    AppendAutoSuggestion(), has_focus(self.buffer) & ~is_done
                ),
                ConditionalProcessor(
                    processor=PasswordProcessor(), filter=to_filter(password)
                ),
                BeforeInput(prompt, style="class:text-area.prompt"),
            ]
            + input_processors,
            search_buffer_control=getattr(self.control, 'search_buffer_control', None),
            preview_search=preview_search,
            focusable=focusable,
            focus_on_click=focus_on_click,
        )
        
        # 重新创建window，使用新的control，并隐藏光标
        self.window = Window(
            height=self.window.height,
            width=self.window.width,
            dont_extend_height=self.window.dont_extend_height,
            dont_extend_width=self.window.dont_extend_width,
            content=self.control,
            style=self.window.style,
            wrap_lines=Condition(lambda: self.wrap_lines),
            left_margins=self.window.left_margins,
            right_margins=self.window.right_margins,
            get_line_prefix=self.window.get_line_prefix,
        )
    
    def append_text(self, text):
        """
        在当前文本末尾添加文本，保持当前滚动位置不变
        """
        if not text or text == "": 
            return

        # 添加新文本
        new_text = self.text + text
        
        if self.auto_scroll_bottom:
            # 自动滚到底部，将光标设置在文本末尾，让Window自动处理滚动
            new_document = Document(new_text, len(new_text))
            self.document = new_document
        else:
            # 保存当前的垂直滚动位置
            current_vertical_scroll = self.window.vertical_scroll
            
            # 保存当前光标位置
            current_cursor_pos = self.buffer.document.cursor_position
            
            # 直接使用buffer.set_document，保持光标位置不变
            new_document = Document(new_text, current_cursor_pos)
            self.buffer.set_document(new_document, bypass_readonly=True)
            
            # 恢复滚动位置
            self.window.vertical_scroll = current_vertical_scroll

    def is_at_bottom(self):
        """
        检查是否已经滚动到底部
        """
        if not self.window.render_info:
            return True
            
        return (
            self.window.render_info.content_height <= 
            self.window.render_info.vertical_scroll + self.window.render_info.window_height
        )
    
    def scroll_to_bottom(self):
        """
        手动滚动到底部
        """
        if self.window.render_info:
            self.window.vertical_scroll = max(
                0, 
                self.window.render_info.content_height - self.window.render_info.window_height
            )